from contextvars import ContextVar

from peewee import _ConnectionState, MySQLDatabase
from playhouse.pool import PooledMySQLDatabase

db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
db_state = ContextVar("db_state", default=db_state_default.copy())

from fastapi import FastAPI
from afast_core.core_plugins.pluginbase import IBasePlugin as BasePlugin
from pydantic import BaseSettings as Settings
from afast_core.core_db.peewee_postgresql import PeeweeSyncPooledPostgresqlClient


class PeeweeConnectionState(_ConnectionState):
    # 连接代理的获取
    def __init__(self, **kwargs):
        super().__setattr__("_state", db_state)
        super().__init__(**kwargs)

    def __setattr__(self, name, value):
        self._state.get()[name] = value

    def __getattr__(self, name):
        return self._state.get()[name]


class PeeweeContextPluginClient(PeeweeSyncPooledPostgresqlClient, BasePlugin):
    '''
    使用示例：
    '''
    # 设置插件默认的参数信息
    name = 'PeeweeContextPluginClient'

    class PeeweeSettings(Settings):
        db_name: str
        db_user: str
        db_pass: str
        db_host: str = '127.0.0.1'
        db_port: int = 5432
        data_base: str = None
        max_connections: int = 60
        stale_timeout: int = 300
        timeout: int = 20

    async def reset_db_state(self):
        self.db._state._state.set(db_state_default.copy())
        self.db._state.reset()

    def setup(self, app: FastAPI, name: str = None, *args, **kwargs):
        """插件初始化"""
        self.app = app
        self.db = None

        @app.on_event("startup")
        def startup_event():
            self.db = MySQLDatabase(
                self.settings.DATABASE,
                user=self.settings.db_user,
                host=self.settings.db_host,
                password=self.settings.db_pass,
                port=self.settings.db_port
            )
            self.db._state = PeeweeConnectionState()
            app.sync_peewee_mysql = self.db._state

        @app.on_event("shutdown")
        def shutdown_event():
            pass
